package com.sunny.zkdislockadvance;

import org.apache.zookeeper.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * @Author zsunny
 * @Date 2018/9/11 10:37
 * @Mail zsunny@yeah.net
 */
public class ZKLock {
    private Logger log = LoggerFactory.getLogger(ZKLock.class);
    private static final String LOCK_ROOT = "/lock";
    private static final String DEFAULT_HOST = "222.201.145.179";
    private static final int DEFAULT_PORT = 5055;
    private CountDownLatch countDownLatch = new CountDownLatch(1);

    private ExecutorService tp = Executors.newFixedThreadPool(5);

    private ZooKeeper zooKeeper;

    public ZKLock(String host, int port) {

        try {
            zooKeeper = new ZooKeeper(host + ":" + port, 20000,
                    watchedEvent -> {
                        log.info("连接完成");
                        try {
                            if(null == zooKeeper.exists(LOCK_ROOT, false)){
                                zooKeeper.create(LOCK_ROOT,"".getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
                            }else{
                                List<String> nodes = zooKeeper.getChildren(LOCK_ROOT, false);
                                nodes.forEach(node->{
                                    try {
                                        zooKeeper.delete(LOCK_ROOT+"/"+node, -1);
                                    } catch (InterruptedException | KeeperException e) {
                                        e.printStackTrace();
                                    }
                                });
                            }
                        } catch (KeeperException e) {
                            e.printStackTrace();
                        } catch (InterruptedException e) {
                            e.printStackTrace();
                        }
                        countDownLatch.countDown();
                    });

            countDownLatch.await();
        } catch (InterruptedException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        }

    }

    public boolean tryLock()
            throws UnknownHostException, KeeperException, InterruptedException {
        CountDownLatch lockWait = new CountDownLatch(1);
        String nodeName = getName();
        //************************************可重入处理
        List<String> nodes = zooKeeper.getChildren(LOCK_ROOT, false);
        Collections.sort(nodes);
        boolean exist = false;
        for(int i=0;i<nodes.size();i++){
            String existData = getData(LOCK_ROOT + "/" + nodes.get(i));
            String[] existDatas = existData.split(";");
            if(nodeName.equals(existDatas[0])){
                if(0 == i){
                    int cnt = Integer.parseInt(existDatas[1]);
                    cnt++;
                    String nodeData = nodeName + ";" + cnt;
                    zooKeeper.setData(LOCK_ROOT + "/" + nodes.get(i), nodeData.getBytes(), -1);
                    return true;
                }else{
                    exist = true;
                }
            }
        }
        //不存在节点则加入
        if(!exist) {
            String nodeData = nodeName + ";0";
            zooKeeper.create(LOCK_ROOT + "/s", nodeData.getBytes(),
                    ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
        }
        //重新获取节点判断
        nodes = zooKeeper.getChildren(LOCK_ROOT, false);
        Collections.sort(nodes);
        System.out.println(nodes);
        String lockData = getData(LOCK_ROOT + "/" + nodes.get(0));
        String[] lockDatas = lockData.split(";");
        if(nodeName.equals(lockDatas[0])){
            System.out.println(nodeName+" "+getData(LOCK_ROOT + "/" + nodes.get(0)));
            return true;
        }
        for(int i=0;i<nodes.size()-1;i++){
            String existData = getData(LOCK_ROOT + "/" + nodes.get(i+1));
            String[] existDatas = existData.split(";");
            if(nodeName.equals(existDatas[0])){
                String prevPath = LOCK_ROOT + "/" + nodes.get(i);
                zooKeeper.exists(prevPath,
                        watchedEvent -> {
                            try {
                                if(null == zooKeeper.exists(prevPath, null) || tryLock()) {
                                    lockWait.countDown();
                                }
                            } catch (KeeperException e) {
                                e.printStackTrace();
                            } catch (InterruptedException e) {
                                e.printStackTrace();
                            } catch (UnknownHostException e) {
                                e.printStackTrace();
                            }
                        });
            }
        }
        System.out.println(getName()+" 阻塞中");
        lockWait.await();
        System.out.println(getName()+" 阻塞完成");
        return true;
    }

    public void unlock() throws KeeperException, InterruptedException, UnknownHostException {
        List<String> nodes = zooKeeper.getChildren(LOCK_ROOT, false);
        Collections.sort(nodes);
        String firstData = getData(LOCK_ROOT + "/" + nodes.get(0));
        String[] firstDatas = firstData.split(";");
        if(!getName().equals(firstDatas[0])){
            return;
        }
        System.out.println(firstDatas[0] + " " + firstDatas[1]);
        int cnt = Integer.parseInt(firstDatas[1]);
        if(0 == cnt){
            zooKeeper.delete(LOCK_ROOT + "/" + nodes.get(0), -1);
        }else {
            cnt--;
            zooKeeper.setData(LOCK_ROOT + "/" + nodes.get(0), (firstDatas[0]+";"+cnt).getBytes(), -1);
        }
    }



    public String getData(String path) throws KeeperException, InterruptedException {
        return new String(zooKeeper.getData(path, false, null));
    }

    public String getName() throws UnknownHostException {
        InetAddress inetAddress = InetAddress.getLocalHost();
        String hostName = inetAddress.getHostName();
        String threadName = Thread.currentThread().getName();
        return hostName+"-"+threadName;
    }

    public static void main(String[] args) {

        ZKLock zkLock = new ZKLock(DEFAULT_HOST, DEFAULT_PORT);


        for(int i=0;i<5;i++)
            zkLock.tp.execute(zkLock.new MyTask());

    }


    class MyTask implements Runnable{
        @Override
        public void run() {
            try {
                if(tryLock()) {
                    System.out.println(Thread.currentThread().getName() + "获得锁");
                    Thread.sleep(10000);
                    System.out.println(Thread.currentThread().getName() + "释放锁");
                    unlock();
                }
            } catch (UnknownHostException | KeeperException | InterruptedException e) {
                e.printStackTrace();
            }
        }
    }

}
